#ifndef UnstructMG_SCAL_KERNERLS_HPP
#define UnstructMG_SCAL_KERNERLS_HPP

#include <cstring>


#ifdef COL_MAJOR
#define IDX(r, c) (c)*dof+(r)
#else // row major
#define IDX(r, c) (r)*dof+(c)
#endif

template<typename data_t, int dof>
inline void d_vec_zero(data_t * dst)
{
    for (int f = 0; f < dof; f++) dst[f] = 0.0;
}

template<typename data_t, int dof>
inline void mat_copy(const data_t * src, data_t * dst)
{
    for (int f = 0; f < dof * dof; f++) dst[f] = src[f];
}

template<typename data_t, int dof>
inline void mat_zero(data_t * dst)
{
    for (int f = 0; f < dof * dof; f++) dst[f] = 0.0;
}

template<typename idx_t, typename data_t, int dof>
inline void set_identity(data_t * vals)
{
    constexpr int e_size = dof*dof;
    for (int f = 0; f < e_size; f++) vals[f] = 0.0;
    for (int f = 0; f <  dof  ; f++) vals[IDX(f, f)] = 1.0;
}

template<typename dst_t, typename src_t, int dof>
inline void mat_trans(dst_t * dst_ptr, const src_t * src_ptr)
{
    for (int r = 0; r < dof; r++)
    for (int c = 0; c < dof; c++)
        dst_ptr[IDX(r, c)] = src_ptr[IDX(c, r)];
}

// out -= mat * in
template<typename idx_t, typename data_t, typename calc_t, int dof>
inline void matvec_mls(const data_t * mat, const calc_t * in, calc_t * out) {
#ifdef COL_MAJOR
    for (int c = 0; c < dof; c++) {
    for (int r = 0; r < dof; r++) {   
        out[r] -= mat[IDX(r, c)] * in[c];
    }}
#else
    assert(false);
#endif
}

// out += mat * in
template<typename idx_t, typename data_t, typename calc_t, int dof>
inline void matvec_mla(const data_t * mat, const calc_t * in, calc_t * out) {
#ifdef COL_MAJOR
    for (int c = 0; c < dof; c++) {
    for (int r = 0; r < dof; r++) {   
        out[r] += mat[IDX(r, c)] * in[c];
    }}
#else
    #pragma GCC unroll (4)
    for (idx_t r = 0; r < dof; r++) {
        calc_t tmp = 0.0;
        const data_t * vals = mat + r * dof;
        #pragma GCC unroll (4)
        for (idx_t c = 0; c < dof; c++)
            tmp += vals[c] * in[c];
        out[r] += tmp;
    }
#endif
}

// out := alpha*mat*in + beta*b
template<typename idx_t, typename data_t, typename calc_t, int dof>
inline void matvec_mla( const calc_t alpha, const data_t * mat, const calc_t * in,
                        const calc_t beta , const calc_t * b, calc_t * out) {
#ifdef COL_MAJOR
    calc_t tmp [dof];
    for (int r = 0; r < dof; r++) tmp[r] = 0.0;
    
    for (int c = 0; c < dof; c++) {
    for (int r = 0; r < dof; r++) {   
        tmp[r] += mat[IDX(r, c)] * in[c];
    }}

    for (int r = 0; r < dof; r++) {
        out[r] = alpha * tmp[r] + beta * b[r];
    }
#else
    #pragma GCC unroll (4)
    for (idx_t r = 0; r < dof; r++) {
        calc_t tmp = 0.0;
        const data_t * vals = mat + r * dof;
        #pragma GCC unroll (4)
        for (idx_t c = 0; c < dof; c++)
            tmp += vals[c] * in[c];
        out[r] = alpha * tmp + beta * b[r];
    }
#endif
}

// out += mat*in*vec
template<typename idx_t, typename data_t, typename calc_t, int dof>
inline void matvec_mla(const data_t * mat, const calc_t * vec, const calc_t * in, calc_t * out) {
#ifdef COL_MAJOR
    for (int c = 0; c < dof; c++) {
    for (int r = 0; r < dof; r++) {   
        out[r] += mat[IDX(r, c)] * in[c] * vec[c];
    }}
#else
    #pragma GCC unroll (4)
    for (idx_t r = 0; r < dof; r++) {
        calc_t tmp = 0.0;
        const data_t * vals = mat + r * dof;
        #pragma GCC unroll (4)
        for (idx_t c = 0; c < dof; c++)
            tmp += vals[c] * in[c] * vec[c];
        out[r] += tmp;
    }
#endif
}

// out := wgt * mat * in
template<typename idx_t, typename data_t, typename calc_t, int dof>
inline void matvec_mul(const data_t * mat, const calc_t * in, calc_t * out, const calc_t wgt) {
#ifdef COL_MAJOR
    calc_t tmp [dof];
    for (int r = 0; r < dof; r++) tmp[r] = 0.0;

    for (int c = 0; c < dof; c++) {
    for (int r = 0; r < dof; r++) {
        tmp[r] += mat[IDX(r, c)] * in[c];
    }}

    for (int r = 0; r < dof; r++) out[r] = wgt * tmp[r];
#else
    #pragma GCC unroll (4)
    for (idx_t r = 0; r < dof; r++) {
        calc_t tmp = 0.0;
        const data_t * vals = mat + r * dof;
        #pragma GCC unroll (4)
        for (idx_t c = 0; c < dof; c++)
            tmp += vals[c] * in[c];
        out[r] = wgt * tmp;
    }
#endif
}

// 原位的乘加
// inout := wgt * mat * inout
template<typename idx_t, typename data_t, typename calc_t, int dof>
inline void matvec_mul(const data_t * mat, calc_t * inout, const calc_t wgt) {
#ifdef COL_MAJOR
    calc_t tmp [dof];
    for (int r = 0; r < dof; r++) tmp[r] = 0.0;

    for (int c = 0; c < dof; c++) {
    for (int r = 0; r < dof; r++) {
        tmp[r] += mat[IDX(r, c)] * inout[c];
    }}

    for (int r = 0; r < dof; r++) inout[r] = wgt * tmp[r];
#else
    calc_t res[dof];
    #pragma GCC unroll (4)
    for (idx_t r = 0; r < dof; r++) {
        calc_t tmp = 0.0;
        const data_t * vals = mat + r * dof;
        #pragma GCC unroll (4)
        for (idx_t c = 0; c < dof; c++)
            tmp += vals[c] * inout[c];
        res[r] = wgt * tmp;
    }

    #pragma GCC unroll (4)
    for (idx_t r = 0; r < dof; r++) inout[r] = res[r];
#endif
}

template<typename idx_t, typename calc_t, int dof>
inline void vecvec_mul(const calc_t * coeff, calc_t * res) {
    #pragma GCC unroll (4)
    for (idx_t r = 0; r < dof; r++)
        res[r] *= coeff[r];
}

// 矩阵求逆，返回的逆矩阵在E中
template<typename data_t, int dof>
inline void matinv_row(data_t * A, data_t * E) {
    // 构造单位阵
    memset(E, 0.0, sizeof(data_t) * dof * dof);
    #pragma GCC unroll (4)
    for (int r = 0; r < dof; r++)
        E[IDX(r, r)] = 1.0;
    // 初等行变换
    for (int i = 0; i < dof; i++) {
        data_t tmp = A[IDX(i, i)];
        #pragma GCC unroll (4)
        for (int j = 0; j < dof; j++) {
            A[IDX(i, j)] /= tmp;
            E[IDX(i, j)] /= tmp;
        }
        // 此时A(i,i)为1，利用这个1消掉上面、下面的行中该列的非零元
        for (int k = 0; k < dof; k++) {
            if (k == i) continue;
            data_t tmp = A[IDX(k, i)];
            #pragma GCC unroll (4)
            for (int j = 0; j < dof; j++) {
                A[IDX(k, j)] -= tmp * A[IDX(i, j)];
                E[IDX(k, j)] -= tmp * E[IDX(i, j)];
            }
        }
    }
}

// C += alpha*A*B + beta*C
template<typename data_t, int dof>
inline void matmat_mla(const data_t alpha, const data_t * A, const data_t * B, const data_t beta, data_t * C)
{
#ifdef COL_MAJOR
    data_t tmp [dof*dof];// A*B
    for (int f = 0; f < dof*dof; f++) tmp[f] = 0.0;

    for (int k = 0; k < dof; k++) {
        for (int j = 0; j < dof; j++) {
        for (int i = 0; i < dof; i++) {
            tmp[IDX(i, j)] += A[IDX(i, k)] * B[IDX(k, j)];
        }}
    }

    for (int f = 0; f < dof*dof; f++) C[f] = alpha * tmp[f] + beta * C[f];
#else
    for (int i = 0; i < dof; i++)
    for (int j = 0; j < dof; j++) {
        data_t tmp = 0.0;
        for (int k = 0; k < dof; k++) {
            tmp += A[i * dof + k] * B[k * dof + j];
        }
        C[i * dof + j] = alpha * tmp + beta * C[i * dof + j];
    }
#endif
}

// C := alpha*A*B
template<typename data_tA, typename data_tB, typename data_tC, int dof>
inline void matmat_mul(const data_tC alpha, const data_tA * A, const data_tB * B, data_tC * C)
{
    if constexpr (dof == 1) {
        C[0] = alpha * A[0] * B[0];
    }
    else {static_assert(dof > 1);
#ifdef COL_MAJOR
        data_tC tmp [dof*dof];// A*B
        for (int f = 0; f < dof*dof; f++) tmp[f] = 0.0;

        for (int k = 0; k < dof; k++) {
            for (int j = 0; j < dof; j++) {
            for (int i = 0; i < dof; i++) {
                tmp[IDX(i, j)] += A[IDX(i, k)] * B[IDX(k, j)];
            }}
        }

        for (int f = 0; f < dof*dof; f++) C[f] = alpha * tmp[f];
#else
        for (int i = 0; i < dof; i++)
        for (int j = 0; j < dof; j++) {
            data_tC tmp = 0.0;
            for (int k = 0; k < dof; k++) {
                tmp += A[i * dof + k] * B[k * dof + j];
            }
            C[i * dof + j] = alpha * tmp;
        }
#endif
    }
}

template<typename data_t, int dof>
inline data_t matNorm_inf(const data_t * vals)
{
    data_t ret = 0.0;
    for (int r = 0; r < dof; r ++) {
        data_t tmp = 0.0;
        for (int c = 0; c < dof; c ++)
            tmp += fabs(vals[IDX(r, c)]);
        if (ret < tmp)
            ret = tmp;
    }
    return ret;
}

template<typename data_t, int dof>
inline data_t matNorm_Frob(const data_t * vals)
{
    data_t ret = 0.0;
    for (int f = 0; f < dof*dof; f++) ret += vals[f]*vals[f];
    return sqrt(ret);
}

#undef IDX

#endif

